from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.utilities import SQLDatabase

from AliUtils import get_model

# 创建模型
model = get_model()

# sqlalchemy 初始化MySQL数据库的连接
HOSTNAME = 'localhost'
PORT = '3306'
DATABASE = 'ry-vue'
USERNAME = 'root'
PASSWORD = 'root'
# mysqlclient驱动URL
# MYSQL_URI = 'mysql+mysqldb://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)

# pymysql驱动URL
MYSQL_URI = 'mysql+pymysql://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)

db = SQLDatabase.from_uri(MYSQL_URI)

# 测试连接是否成功
# print(db.get_usable_table_names())
# print(db.run('select * from sys_user limit 10;'))

# 直接使用大模型和数据库整合, 只能根据你的问题生成SQL
# 初始化生成SQL的chain
test_chain = create_sql_query_chain(model, db)
resp = test_chain.invoke({'question': '请问：用户表中有多少条数据？'})
print(resp)
